# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     test
   Description :   
   Author :       lth
   date：          2023/1/30
-------------------------------------------------
   Change Activity:
                   2023/1/30 10:29: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import torch
from torch.backends import cudnn
from transformers import BertTokenizer

from config import GetConfig
from model import BertClassifier

tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
labels = {
    0: "business",
    1: "entertainment",
    2: "sport",
    3: "tech",
    4: "politics"
}


class Test:
    def __init__(self):
        self.args = GetConfig()
        print(f"-----------{self.args.project_name}-------------")
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if use_cuda else "cpu")

        self.model = BertClassifier(type="test").to(self.device)

        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))
            cudnn.benchmark = True
            cudnn.enabled = True


        self.model.load_state_dict(torch.load("weights/best.pth")["model_state_dict"])


        self.model.eval()

    @torch.no_grad()
    def inference(self, x):
        data = tokenizer(x,
                         padding='max_length',
                         max_length=512,
                         truncation=True,
                         return_tensors="pt")

        mask = data["attention_mask"].to(self.device)
        input_id = data["input_ids"].squeeze(1).to(self.device)
        output = self.model(input_id, mask)

        output = output.argmax(dim=1).item()
        return output


if __name__ == "__main__":
    model = Test()
    news = "wales silent on grand slam talk rhys williams says wales are still not thinking of winning the grand slam despite a third six nations win.   that s the last thing on our minds at the moment   said williams  a second- half replacement in saturday s 24-18 win over france in paris.  we all realise how difficult a task it is to go up to scotland and beat them.  we ve come unstuck there a couple of times recently so our focus is on that game and we ll worry about ireland hopefully after we ve beaten scotland.  with captain gareth thomas ruled out of the rest of the campaign with a broken thumb  williams is vying for his first start in the championship so far.  kevin morgan is probably favourite to replace thomas at full-back  leaving williams and hal luscombe to battle for the right wing berth. a hamstring injury denied luscombe the opportunity to make a third successive start  but the dragons winger is expected to be fit for the trip to murrayfield on 13 march. hooker robin mcbryde is doubtful after picking up a knee injury in paris  but centre sonny parker and flanker colin charvis are set to recover from injury to be in contention for selection. said wales assistant coach scott johnson:  they ve worked through the weekend and the reports are a bit more positive.  so we re getting a couple back and that adds to the depth of the squad.  scotland secured their first win of the campaign on saturday by grinding out an 18-10 win over italy.  matt williams  side has shown little in attack  but johnson insisted the scots will be difficult opposition to break down.  italy are really brave opposition and sometimes it s very hard to win   he said.  so an ugly win can be just as effective as a 30 or 40 point victory.  scotland are a hard side and very underrated so we re not taking anything for granted.  we re not basking in the glory of winning our first three games. we ve got to be diligent in our preparation.  that s my job and we ve got to make sure we re focused."

    result = model.inference(news)
    print(labels[result])
